# coding=UTF-8
# author=suemi
# created at 16/5/13
import requests
from elasticsearch import Elasticsearch
from utils import Constants


class QueryBuilder(object):
    """
    从用户输入中生成符合ES要求的请求体并解析ES返回的响应
    """

    address = Constants.ES_SETTINGS['host'] + ':' + str(Constants.ES_SETTINGS['port'])
    es = {address: Elasticsearch([address])}

    def __init__(self,
                 indexName=Constants.INDEX_NAME,
                 indexType=Constants.INDEX_TYPE):
        self.indexName = indexName
        self.indexType = indexType
        self.toCluster = False
        self.start = 0
        self.size = None
        self.sortBy = ['_score']
        self.matchFields = {}
        self.word = None
        self.clusterAlgorithm = 'lingo'
        self.address = self.__class__.address

    @staticmethod
    def instance(indexName=Constants.INDEX_NAME,
                 indexType=Constants.INDEX_TYPE):
        return QueryBuilder(indexName, indexType)

    def connect(self):
        """
        :return:获取ES连接实例
        """
        tmp = self.__class__.es
        if not tmp.has_key(self.address):
            tmp[self.address] = Elasticsearch([self.address])
        return tmp.get(self.address)

    def specify(self,
                indexName=Constants.INDEX_NAME,
                indexType=Constants.INDEX_TYPE):
        """
        指定索引名和索引类型
        :param indexName:
        :param indexType:
        :return:
        """
        self.indexName = indexName
        self.indexType = indexType
        return self

    def term(self, word, field='_all'):
        """
        添加查询词
        :param word:
        :param field:
        :return:
        """
        self.matchFields[field] = word
        self.word = word
        return self

    def locate(self, host, port=9200):
        """
        修改服务器地址
        :param host:
        :param port:
        :return:
        """
        self.address = host + ':' + str(port)
        return self

    def cluster(self, toCluster):
        """
        设定是否需要对搜索结果聚类
        :param toCluster:
        :return:
        """
        self.toCluster = toCluster
        return self

    def algorithm(self, clusterAlgorithm):
        """
        设定聚类算法
        :param clusterAlgorithm:
        :return:
        """
        if clusterAlgorithm in ['lingo', 'stc', 'kmeans', 'byurl']:
            self.clusterAlgorithm = clusterAlgorithm
        return self

    def limit(self, size):
        """
        返回返回结果上限
        :param size:
        :return:
        """
        self.size = size
        return self

    def sort(self, sortItem):
        """
        设定返回文档的排序策略
        :param sortItem:
        :return:
        """
        tmp = sortItem.split(':')
        if len(tmp) == 1:
            self.sortBy.insert(-1, tmp[0])
        elif len(tmp) == 2:
            self.sortBy.insert(-1, {tmp[0]: {'order': tmp[1]}})
        else:
            raise ValueError
        return self

    def cursor(self, start):
        """
        指定返回文档的起始位置,便于分页
        :param start:
        :return:
        """
        self.start = start
        return self

    def submit(self, document):
        """
        提交一个文档
        :param document:
        :return:
        """
        return self.connect().index(self.indexName, self.indexType, document.toDict(), id=document.documentId())

    def delete(self, document):
        """
        删除一个文档
        :param document:
        :return:
        """
        return self.connect().delete(self.indexName, self.indexType, document.documentId())

    def persist(self, document, delta=None):
        """
        更新或插入一个文档
        :param document:
        :param delta:
        :return:
        """
        if self.connect().exists(self.indexName, self.indexType, document.documentId()):
            return self.connect().update(self.indexName, self.indexType, document.documentId(),
                                         document.toDict() if delta is None else delta)
        else:
            return self.submit(document)

    def load(self, documentId):
        """
        根据Id获取文档
        :param documentId:
        :return:
        """
        return self.connect().get_source(self.indexName, self.indexType, documentId)

    # def query(self):
    #     """
    #     使用Pyes提供的Python API来查询
    #     :return:
    #     """
    #
    #     if not self.toCluster:
    #         if len(self.matchFields) == 0:
    #             # generate_model(self.indexName,self.indexType).objects.filter
    #             q = MatchAllQuery()
    #         else:
    #             item = self.matchFields.items()[-1]
    #             q = MatchQuery(item[0], item[1])
    #         q = Search(q, start=self.start, size=self.size)
    #         sortBy = ''
    #         for item in self.sortBy:
    #             if isinstance(item, str):
    #                 sortBy += item + ','
    #             elif isinstance(item, dict):
    #                 tmp = item.popitem()
    #                 sortBy += tmp[0]
    #                 sortBy += ':' + tmp[1]['order']
    #                 sortBy += ','
    #         sortBy = sortBy[:-1] if len(sortBy) > 0 else sortBy
    #         response = self.connect().search(q, indexes=[self.indexName], sort=sortBy)
    #         items = [item for item in response]
    #         orignal = response._results
    #         for i in range(len(items)):
    #             tmp = dict.fromkeys(items[i].keys(), '')
    #             for key in tmp.keys():
    #                 tmp[key] = items[i][key]
    #             orignal['hits']['hits'][i]['_source'] = tmp
    #
    #         return self.parseResult(orignal)
    #     else:
    #         return self.rawQuery()

    def query(self):
        """
        直接使用ES的Restful请求来查询
        :return:
        """
        query = {'match_all': {}} if len(self.matchFields) == 0 else {
            'match': self.matchFields
        }
        searchRequest = {
            'query': query,
            "sort": self.sortBy,
            'from': self.start,
            'highlight': {
                'pre_tags': ['<strong>'],
                'post_tags': ['</strong>'],
                # 此处非常重要,由于我们设置的查找域基本为_all,不设置的话就会导致任一域都不高亮
                'require_field_match': False,
                'fields': {
                    'title': {
                        # 'type': 'fvh', 去掉注释返回高亮的偏移位置
                        'fragment_size': '20',
                        "number_of_fragments": 1

                    },
                    'content': {
                        # 'type': 'fvh',
                        'fragment_size': '100',
                        "number_of_fragments": 3
                    }
                }
            }
        }
        if self.size is not None:
            searchRequest['size'] = self.size

        url = 'http://' + self.address + '/' + self.indexName + '/' + self.indexType
        if self.toCluster:
            requestBody = {
                'search_request': searchRequest,
                'query_hint': self.word,
                'field_mapping': {
                    'title': ['_source.title'],
                    'content': ['_source.content'],
                    'lang': ['_source.lang']
                },
                'algorithm': self.clusterAlgorithm,
            }
            url += '/_search_with_clusters'
            response = requests.post(url, json=requestBody).json()
        else:
            url += '/_search'
            response = self.connect().search(index=self.indexName, doc_type=self.indexType, body=searchRequest)
        return self.parseResult(response)

    def parseResult(self, response):
        """
        解析ES返回的结果
        :param response:
        :return:
        """

        result = {
            'abstract': {},
            'documents': [],
            'detail': {},
            'success': True
        }
        if response.has_key('error'):
            result['success'] = False
            result['cause'] = response['error']['root_cause']
            return result

        hits = response.get('hits', {
            'total': 0,
            'max_score': 0,
            'hits': []
        })
        result['abstract']['total'] = hits['total']
        result['abstract']['max_score'] = hits['max_score']
        id2index = {}
        for i in range(len(hits['hits'])):
            hit = hits['hits'][i]
            tmp = hit['_source']
            tmp['_score'] = hit['_score']
            tmp['_id'] = hit['_id']
            tmp['highlight'] = hit.get("highlight", {})
            result['documents'].append(tmp)
            id2index[hit['_id']] = i

        if self.toCluster:
            clusters = response.get('clusters', [])
            result['abstract']['cluster_num'] = len(clusters)
            result['detail']['clusters'] = clusters
            for cluster in clusters:
                for documentId in cluster['documents']:
                    position = id2index[documentId]
                    document = result['documents'][position]
                    document['label'] = cluster['label']

        return result
